#include "prolong.hpp"
#include "scal_kernels.hpp"

template<typename idx_t, typename data_t, int dof>
par_CSRMatrix<idx_t, data_t, data_t, dof>* Jacobi_A(const data_t weight,
    const par_CSRMatrix<idx_t, data_t, data_t, dof> & A)
{
    par_CSRMatrix<idx_t, data_t, data_t, dof> * jac_mat
        = new par_CSRMatrix<idx_t, data_t, data_t, dof>(A.comm, A.rows_partition, A.cols_partition);

    jac_mat->diag.nrows = A.diag.nrows;
    jac_mat->diag.ncols = A.diag.ncols;
    jac_mat->diag.nnz   = A.diag.nnz  ; jac_mat->diag.alloc_mem(false);

    jac_mat->offd.nrows = A.offd.nrows;
    jac_mat->offd.ncols = A.offd.ncols;
    jac_mat->offd.nnz   = A.offd.nnz  ; jac_mat->offd.alloc_mem(false);

    jac_mat->col_map_offd = A.col_map_offd;
    jac_mat->commpkg      = A.commpkg;

    const idx_t nrows = jac_mat->end_row - jac_mat->beg_row;
    constexpr int e_size = dof*dof;

    data_t eye [e_size];
    for (int f = 0; f < e_size; f++) eye[f] = 0.0;
    for (int f = 0; f < dof; f++) eye[f * dof + f] = 1.0;// 对角元为1

    for (idx_t i = 0; i < nrows; i++) {
        data_t invD [e_size];
        for (idx_t p = A.diag.row_ptr[i]; p < A.diag.row_ptr[i + 1]; p++) {
            if (A.diag.col_idx[p] == i) {// 找到了对角元
                data_t ibuf [e_size];
                for (int f = 0; f < e_size; f++)
                    ibuf[f] = (A.diag.vals + p * e_size)[f];
                matinv_row<data_t, dof>(ibuf, invD);
                break;
            }
        }

        for (int f = 0; f < e_size; f++) invD[f] *= - weight;// 乘以权重，此时invD中保存 -w*D^{-1}

        jac_mat->diag.row_ptr[i] = A.diag.row_ptr[i];
        for (idx_t p = A.diag.row_ptr[i]; p < A.diag.row_ptr[i + 1]; p++) {
            jac_mat->diag.col_idx[p] = A.diag.col_idx[p];
            matmat_mul<data_t, data_t, data_t, dof>(1.0, invD, A.diag.vals + p * e_size,
                                                        jac_mat->diag.vals + p * e_size);// -w*D{-1}*A
            if (A.diag.col_idx[p] == i) {// 是对角元，需要类加上单位矩阵
                for (int f = 0; f < e_size; f++)
                    (jac_mat->diag.vals + p * e_size)[f] += eye[f];// I -w*D{-1}*A
            }
        }

        jac_mat->offd.row_ptr[i] = A.offd.row_ptr[i];
        for (idx_t p = A.offd.row_ptr[i]; p < A.offd.row_ptr[i + 1]; p++) {
            jac_mat->offd.col_idx[p] = A.offd.col_idx[p];
            matmat_mul<data_t, data_t, data_t, dof>(1.0, invD, A.offd.vals + p * e_size,
                                                        jac_mat->offd.vals + p * e_size);// -w*D{-1}*A
        }
    }
    jac_mat->diag.row_ptr[nrows] = A.diag.row_ptr[nrows];
    jac_mat->offd.row_ptr[nrows] = A.offd.row_ptr[nrows];

    return jac_mat;
}

template par_CSRMatrix<int, float , float , 1>* Jacobi_A(const float , const par_CSRMatrix_I32F32F32_DOF1 &);
template par_CSRMatrix<int, float , float , 2>* Jacobi_A(const float , const par_CSRMatrix_I32F32F32_DOF2 &);
template par_CSRMatrix<int, float , float , 3>* Jacobi_A(const float , const par_CSRMatrix_I32F32F32_DOF3 &);
template par_CSRMatrix<int, float , float , 4>* Jacobi_A(const float , const par_CSRMatrix_I32F32F32_DOF4 &);

template par_CSRMatrix<int, double, double, 1>* Jacobi_A(const double, const par_CSRMatrix_I32F64F64_DOF1 &);
template par_CSRMatrix<int, double, double, 2>* Jacobi_A(const double, const par_CSRMatrix_I32F64F64_DOF2 &);
template par_CSRMatrix<int, double, double, 3>* Jacobi_A(const double, const par_CSRMatrix_I32F64F64_DOF3 &);
template par_CSRMatrix<int, double, double, 4>* Jacobi_A(const double, const par_CSRMatrix_I32F64F64_DOF4 &);


